import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision import datasets
import torch.optim as optim

transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.1307,), (0.3081,))])

train_dataset = datasets.MNIST(root='../dataset/mnist',
                               train=True,
                               transform=transform,
                               download=True)
test_dataset = datasets.MNIST(root='../dataset/mnist',
                              train=False,
                              transform=transform,
                              download=True)
train_dataloader = DataLoader(dataset=train_dataset,
                              batch_size=64,
                              shuffle=True)
test_dataloader = DataLoader(dataset=test_dataset,
                             batch_size=64,
                             shuffle=False)


class InceptionA(torch.nn.Module):
    def __init__(self, in_channels):
        super(InceptionA, self).__init__()
        # The first style of 1x1 kernel
        self.Conv1x1_1 = torch.nn.Conv2d(in_channels, 24, kernel_size=(1, 1))
        self.Conv1x1_2 = torch.nn.Conv2d(in_channels, 16, kernel_size=(1, 1))
        # 3x3 kernel
        self.Conv3x3_1 = torch.nn.Conv2d(16, 24, kernel_size=(3, 3), padding=1)
        self.Conv3x3_2 = torch.nn.Conv2d(24, 24, kernel_size=(3, 3), padding=1)
        # 5x5 kernel
        self.Conv5x5_1 = torch.nn.Conv2d(16, 24, kernel_size=(5, 5), padding=2)

    def forward(self, x):
        # Average pooling branch
        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1)
        branch_pool = self.Conv1x1_1(branch_pool)
        # 1x1 kernel branch
        branch_1x1 = self.Conv1x1_2(x)
        # 5x5 kernel branch
        branch_5x5 = self.Conv1x1_2(x)
        branch_5x5 = self.Conv5x5_1(branch_5x5)
        # 3x3 kernel branch
        branch_3x3 = self.Conv1x1_2(x)
        branch_3x3 = self.Conv3x3_1(branch_3x3)
        branch_3x3 = self.Conv3x3_2(branch_3x3)
        outputs = [branch_pool, branch_1x1, branch_3x3, branch_5x5]
        return torch.cat(outputs, dim=1)


class GoogleNet(torch.nn.Module):
    def __init__(self):
        super(GoogleNet, self).__init__()
        self.Conv_1 = torch.nn.Conv2d(1, 10, kernel_size=(5, 5))
        self.Conv_2 = torch.nn.Conv2d(88, 20, kernel_size=(5, 5))
        self.incep_1 = InceptionA(in_channels=10)
        self.incep_2 = InceptionA(in_channels=20)
        self.mp = torch.nn.MaxPool2d(2)
        self.fc = torch.nn.Linear(1408, 10)

    def forward(self, x):
        batch_size = x.size(0)
        x = F.relu(self.mp(self.Conv_1(x)))
        x = self.incep_1(x)
        x = F.relu(self.mp(self.Conv_2(x)))
        x = self.incep_2(x)
        x = x.view(batch_size, -1)
        x = self.fc(x)
        return x


googleNet = GoogleNet()
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
googleNet.to(device)
print(torch.cuda.is_available())
googleNet.to(device)
criterion = torch.nn.CrossEntropyLoss()
optimizer = optim.SGD(googleNet.parameters(), lr=0.01, momentum=0.5)


# Set train cycle
def train(epoch):
    running_loss = 0.0
    for batch_idx, data in enumerate(train_dataloader, 0):
        inputs, target = data
        # Set GPU
        inputs, target = inputs.to(device), target.to(device)
        optimizer.zero_grad()
        # Forward + Backward + Update
        outputs = googleNet(inputs)
        loss = criterion(outputs, target)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if batch_idx % 300 == 299:
            print('[%d,%5d] loss: %.6f' % (epoch + 1, batch_idx + 1, running_loss / 300))
            running_loss = 0.0


def test():
    correct = 0
    total = 0
    with torch.no_grad():
        for data in test_dataloader:
            images, labels = data
            images, labels = images.to(device), labels.to(device)
            outputs = googleNet(images)
            _, predicted = torch.max(outputs.data, dim=1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    print('Accuracy on test set : %.3f' % (100 * correct / total))


if __name__ == "__main__":
    for epoch in range(30):
        train(epoch)
        test()
